import concurrent.futures
import os
import subprocess
import sys
from datetime import datetime


PROJECT_ROOT = "/home/ma-user/work/MultiKernelBench"
EVAL_SCRIPT = os.path.join(PROJECT_ROOT, "eval_single_runner.py")
ADD_PROMPT_FILE = os.path.join(PROJECT_ROOT, "prompts", "tilelang_new_model_add.py")
GEMM_PROMPT_FILE = os.path.join(PROJECT_ROOT, "prompts", "tilelang_new_model_gemm.py")


def build_command(eval_type: str) -> list:
    python_executable = sys.executable or "python"
    if eval_type == "add":
        prompt_file = ADD_PROMPT_FILE
        operation = "add"
        category = "math"
    elif eval_type == "gemm":
        prompt_file = GEMM_PROMPT_FILE
        operation = "gemm"
        category = "matmul"
    else:
        raise ValueError(f"Unknown eval_type: {eval_type}")
    
    return [
        python_executable,
        EVAL_SCRIPT,
        prompt_file,
        operation,
        "tilelang",
        "result.json",
        category,
    ]


def run_single(index: int, eval_type: str) -> tuple:
    command = build_command(eval_type)
    start_time = datetime.now()
    try:
        completed = subprocess.run(
            command,
            cwd=PROJECT_ROOT,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        duration_s = (datetime.now() - start_time).total_seconds()
        return (
            index,
            eval_type,
            completed.returncode,
            duration_s,
            completed.stdout,
            completed.stderr,
        )
    except Exception as exc:  # noqa: BLE001
        duration_s = (datetime.now() - start_time).total_seconds()
        return (index, eval_type, -9999, duration_s, "", f"Exception: {exc}")


def main(total_runs: int = 10, max_workers: int = 10) -> int:
    if max_workers is None:
        max_workers = total_runs

    # Create tasks for both add and gemm evaluations
    eval_types = ["add", "gemm"]
    total_tasks = total_runs * len(eval_types)
    
    print(
        f"Starting {total_tasks} parallel runs ({total_runs} runs each for add and gemm, max_workers={max_workers}) at {datetime.now().isoformat()}"
    )

    results: list[tuple] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit tasks for both evaluation types
        future_to_task = {}
        task_id = 0
        for eval_type in eval_types:
            for i in range(total_runs):
                future = executor.submit(run_single, i, eval_type)
                future_to_task[future] = (task_id, eval_type, i)
                task_id += 1
        
        for future in concurrent.futures.as_completed(future_to_task):
            task_id, eval_type, run_idx = future_to_task[future]
            try:
                res = future.result()
                print(res)
                results.append(res)
                _, _, returncode, duration_s, _, _ = res
                status = "OK" if returncode == 0 else f"FAIL({returncode})"
                print(f"Run #{run_idx} ({eval_type}) finished: {status} in {duration_s:.2f}s")
            except Exception as exc:  # noqa: BLE001
                print(f"Run #{run_idx} ({eval_type}) raised exception: {exc}")

    successes = sum(1 for r in results if r[2] == 0)
    failures = len(results) - successes
    print(
        f"All runs completed. Success: {successes}, Failures: {failures} at {datetime.now().isoformat()}"
    )

    # Optionally dump brief stderr for failures to help debugging
    for idx, eval_type, returncode, duration_s, _out, err in sorted(results, key=lambda x: x[0]):
        if returncode != 0:
            print("-" * 60)
            print(f"Run #{idx} ({eval_type}) FAILED (code {returncode}, {duration_s:.2f}s). Stderr:")
            print(err.strip())

    return 0 if failures == 0 else 1


if __name__ == "__main__":
    # Allow overrides via environment variables if desired
    try:
        total = int(os.environ.get("TOTAL_RUNS", "5"))
    except ValueError:
        total = 5
    try:
        workers_env = os.environ.get("MAX_WORKERS")
        workers = int(workers_env) if workers_env else None
    except ValueError:
        workers = None
    sys.exit(main(total_runs=total, max_workers=workers))


